import os, sys, time, re
import glob
import math
import random
import copy
from typing import *
import enum
import itertools
import json
from tqdm import tqdm

import pandas as pd
import numpy as np

import mindspore as ms
import mindspore.nn as nn
import mindspore.ops as ops
import mindspore.numpy as mnp
from mindspore.communication import get_group_size, get_rank, init

from wputils.utils.resize import m3d, m3dc, m3dc_inv, r3d, r3d0
from wputils.utils.norm import nml
from wputils.utils.image import rbbox, bshape, bcenter
# from wpdata.dloader.utils import closest_divisible, one_hot

from wputils.utils.io import rnii, wnii, rnpz, wnpy, rnpy, cdirs, fdirs, gsiz, gmdf, wseq, rseq



def one_hot(arry, num_classes=1):
    """
    Perform one-hot encoding for a numpy array.
    Args:
        arry (numpy.ndarray): Input numpy array
        num_classes (int): Number of unique classes
    Returns:
        numpy.ndarray: One-hot encoded matrix
    """

    return np.eye(num_classes)[arry].transpose(3, 0, 1, 2)


class HematomaDataset(object):
    def __init__(self, path):
        # if isinstance(kind, HematomaType):
        #     kind = kind.value  # cast enum to values

        self.data = fdirs(path, '.npy')[0:16]

    def __len__(self) -> int:
        return len(self.data)

    def __getitem__(self, item):
        elem = self.data[item]
        img_arr = rnpy(os.path.join(elem, 'img.npy'))
        lab_arr = rnpy(os.path.join(elem, 'lab.npy'))

        img_arr = nml(img_arr, 0, 1024 + 2048)
        lab_arr = one_hot(lab_arr, 2)

        return img_arr[np.newaxis].astype(np.float32), lab_arr.astype(np.float32)


if __name__ == '__main__':
    #  sorting files:

    ROOT = '/share_data/liupan/hematoma/npyfiles/24h/train'
    # files = [os.path.join(ROOT, f) for f in os.listdir(ROOT)]
    train_ds = HematomaDataset(ROOT)
    train_dsl = ms.dataset.GeneratorDataset(train_ds, column_names=['img', 'lab'], shuffle=True)
    train_dsl = train_dsl.batch(2, drop_remainder=True)
    train_dsl = train_dsl.repeat(1)
    for ts in train_dsl.create_dict_iterator():
        print('img', ts['img'].shape)
        print('lab', ts['lab'].shape)
